# 从URDF文件输出model，data以及meshcat的viz函数
import pinocchio as pin
from os.path import dirname
import sys
from pinocchio.visualize import MeshcatVisualizer

def robot_viz(URDF_FILE_PATH):
    # 当URDF文件内使用相对路径时，一定要设置base path
    # pinocchio.buildModelsFromUrdf会返回三个模型
    # model, collision_model, visual_model = pinocchio.buildModelsFromUrdf(URDF_FILE_PATH, [os.path.dirname(URDF_FILE_PATH)])
    try:
        robot = pin.RobotWrapper.BuildFromURDF(URDF_FILE_PATH, [dirname(URDF_FILE_PATH)])
    except:
        print("Error while loading the URDF file")
        print(f"URDF_FILE_PATH is: {URDF_FILE_PATH}")
        sys.exit(0)

    try:
        viz = MeshcatVisualizer(robot.model, robot.collision_model, robot.visual_model)
        viz.initViewer(open=True, loadModel=True)
    except ImportError as err:
        print(
            "Error while initializing the viewer. It seems you should install Python meshcat"
        )
        print(err)
        sys.exit(0)

    viz.loadViewerModel()

    # 此处data应从viz获得，这样后期绘制关节速度frame才显示
    # data = viz.data

    return robot, viz